from torch.utils.data import Dataset
import torch


class NewsRatingDataset(Dataset):

    def __init__(self, df):
        self.df = df

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # user data
        uid = self.df.iloc[idx, 0]
        gender = self.df.iloc[idx, 2]
        age = self.df.iloc[idx, 3]
        job = self.df.iloc[idx, 4]

        # news data
        mid = self.df.iloc[idx, 1]
        mtitle = self.df.iloc[idx, 5]
        mtype = self.df.iloc[idx, 6]

        rank = torch.FloatTensor([self.df.iloc[idx, 7]])

        user_input = {
            'uid': torch.LongTensor([uid]),
            'gender': torch.LongTensor([gender]),
            'age': torch.LongTensor([age]),
            'job': torch.LongTensor([job]),
        }

        news_input = {
            'mid': torch.LongTensor([mid]),
            'mtitle': torch.LongTensor([mtitle]),
            'mtype': torch.LongTensor([mtype]).view(-1),
        }

        sample = {
            'user_input': user_input,
            'news_input': news_input,
            'target': rank
        }

        return sample
